"""
AKShare数据获取模块
基于PRD精简版的数据获取和缓存策略
"""

import akshare as ak
import pandas as pd
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
import time
import logging

from config import API_CONFIG, CACHE_CONFIG, VALIDATION_CONFIG
from database import db_manager
from account_mapping import account_classifier

logger = logging.getLogger(__name__)


class DataFetcher:
    """财务数据获取器"""
    
    def __init__(self):
        """初始化数据获取器"""
        self.api_timeout = API_CONFIG['timeout']
        self.retry_times = API_CONFIG['retry_times']
        self.retry_delay = API_CONFIG['retry_delay']
    
    def get_balance_sheet_data(self, stock_code: str, 
                              force_refresh: bool = False) -> pd.DataFrame:
        """
        获取资产负债表数据
        
        Args:
            stock_code: 股票代码（6位数字）
            force_refresh: 是否强制刷新缓存
            
        Returns:
            包含分类信息的资产负债表DataFrame
        """
        # 标准化股票代码
        formatted_code = self._format_stock_code(stock_code)
        
        # 检查缓存（如果不强制刷新）
        if not force_refresh:
            cached_data = self._get_cached_data(formatted_code)
            if cached_data is not None and not cached_data.empty:
                logger.info(f"使用缓存数据: {formatted_code}")
                return cached_data
        
        # 从API获取数据
        logger.info(f"从API获取数据: {formatted_code}")
        raw_data = self._fetch_from_api(formatted_code)
        
        if raw_data is None or raw_data.empty:
            logger.error(f"无法获取股票 {formatted_code} 的数据")
            return pd.DataFrame()
        
        # 数据清洗和验证
        cleaned_data = self._clean_and_validate(raw_data)
        
        # 科目分类
        categorized_data = account_classifier.classify_dataframe(cleaned_data)
        
        # 保存到缓存
        self._save_to_cache(formatted_code, categorized_data)
        
        return categorized_data
    
    def _format_stock_code(self, stock_code: str) -> str:
        """
        格式化股票代码为AKShare要求的格式
        
        Args:
            stock_code: 原始股票代码
            
        Returns:
            格式化后的股票代码（如SH600519）
        """
        # 移除所有非数字字符
        code = ''.join(filter(str.isdigit, stock_code))
        
        if len(code) != 6:
            raise ValueError(f"股票代码必须是6位数字，当前输入: {stock_code}")
        
        # 根据代码判断市场
        if code.startswith(('60', '68')):
            return f"SH{code}"  # 上海证券交易所
        elif code.startswith(('00', '30')):
            return f"SZ{code}"  # 深圳证券交易所
        else:
            # 默认上海
            return f"SH{code}"
    
    def _get_cached_data(self, stock_code: str) -> Optional[pd.DataFrame]:
        """从缓存获取数据"""
        try:
            cached_data = db_manager.get_balance_sheet_data(stock_code)
            
            if not cached_data.empty:
                # 检查数据是否过期
                latest_date = pd.to_datetime(cached_data['created_at']).max()
                cache_age = (datetime.now() - latest_date).days
                
                # 根据报告类型判断缓存是否过期
                if cache_age <= CACHE_CONFIG['annual_report_days']:
                    return self._format_cached_data(cached_data)
            
        except Exception as e:
            logger.error(f"获取缓存数据失败: {e}")
        
        return None
    
    def _format_cached_data(self, cached_data: pd.DataFrame) -> pd.DataFrame:
        """格式化缓存数据为标准格式"""
        if cached_data.empty:
            return cached_data
        
        # 重命名列以匹配API数据格式
        formatted_data = cached_data.rename(columns={
            'report_date': '报告日期',
            'account_name': '科目',
            'amount': '金额',
            'report_type': '报告类型'
        })
        
        return formatted_data
    
    def _fetch_from_api(self, stock_code: str) -> Optional[pd.DataFrame]:
        """从AKShare API获取数据"""
        for attempt in range(self.retry_times):
            try:
                logger.info(f"API调用尝试 {attempt + 1}/{self.retry_times}: {stock_code}")
                
                # 调用AKShare接口
                data = ak.stock_balance_sheet_by_report_em(symbol=stock_code)
                
                if data is not None and not data.empty:
                    logger.info(f"成功获取数据: {stock_code}, 记录数: {len(data)}")
                    return data
                else:
                    logger.warning(f"API返回空数据: {stock_code}")
                    
            except Exception as e:
                logger.error(f"API调用失败 (尝试 {attempt + 1}): {e}")
                
                if attempt < self.retry_times - 1:
                    time.sleep(self.retry_delay)
        
        return None
    
    def _clean_and_validate(self, data: pd.DataFrame) -> pd.DataFrame:
        """数据清洗和验证"""
        if data.empty:
            return data

        # 复制数据避免修改原始数据
        raw_data = data.copy()

        # AKShare返回的是宽表格式，需要转换为长表格式
        # 检查是否有REPORT_DATE列
        if 'REPORT_DATE' not in raw_data.columns:
            logger.error("缺少REPORT_DATE列")
            return pd.DataFrame()

        try:
            # 转换宽表为长表格式
            cleaned_data = self._transform_wide_to_long(raw_data)

            if cleaned_data.empty:
                logger.error("数据转换后为空")
                return pd.DataFrame()

            # 数据类型转换
            cleaned_data['报告日期'] = pd.to_datetime(cleaned_data['报告日期'])
            cleaned_data['金额'] = pd.to_numeric(cleaned_data['金额'], errors='coerce')

            # 移除金额为空或0的行
            cleaned_data = cleaned_data.dropna(subset=['金额'])
            cleaned_data = cleaned_data[cleaned_data['金额'] != 0]

            # 移除重复记录
            before_count = len(cleaned_data)
            cleaned_data = cleaned_data.drop_duplicates(subset=['报告日期', '科目'])
            after_count = len(cleaned_data)

            if before_count != after_count:
                logger.info(f"移除重复记录: {before_count - after_count} 条")

            # 数据范围验证
            min_years = VALIDATION_CONFIG['min_data_years']
            max_years = VALIDATION_CONFIG['max_data_years']

            # 获取最近的数据
            latest_date = cleaned_data['报告日期'].max()
            earliest_date = latest_date - timedelta(days=max_years * 365)

            cleaned_data = cleaned_data[cleaned_data['报告日期'] >= earliest_date]

            # 检查数据年份是否足够
            unique_years = cleaned_data['报告日期'].dt.year.nunique()
            if unique_years < min_years:
                logger.warning(f"数据年份不足: {unique_years} < {min_years}")

            logger.info(f"数据清洗完成，保留 {len(cleaned_data)} 条记录")
            return cleaned_data

        except Exception as e:
            logger.error(f"数据清洗失败: {e}")
            return pd.DataFrame()

    def _transform_wide_to_long(self, data: pd.DataFrame) -> pd.DataFrame:
        """将AKShare的宽表格式转换为长表格式"""
        try:
            # 基础信息列
            base_columns = ['REPORT_DATE', 'REPORT_TYPE', 'SECURITY_NAME_ABBR']

            # 财务科目列（排除基础信息列和YOY列）
            financial_columns = [col for col in data.columns
                               if col not in base_columns
                               and not col.endswith('_YOY')
                               and col not in ['SECUCODE', 'SECURITY_CODE', 'ORG_CODE', 'ORG_TYPE',
                                             'REPORT_DATE_NAME', 'SECURITY_TYPE_CODE', 'NOTICE_DATE',
                                             'UPDATE_DATE', 'CURRENCY', 'OPINION_TYPE', 'OSOPINION_TYPE',
                                             'LISTING_STATE']]

            # 创建长表格式的数据列表
            long_data = []

            for _, row in data.iterrows():
                report_date = row['REPORT_DATE']
                report_type = row.get('REPORT_TYPE', '年报')

                for col in financial_columns:
                    amount = row[col]
                    if pd.notna(amount) and amount != 0:
                        # 将英文科目名转换为中文（简化映射）
                        chinese_name = self._map_english_to_chinese(col)

                        long_data.append({
                            '报告日期': report_date,
                            '报告类型': report_type,
                            '科目': chinese_name,
                            '金额': amount
                        })

            result_df = pd.DataFrame(long_data)
            logger.info(f"宽表转长表完成，转换了 {len(result_df)} 条记录")
            return result_df

        except Exception as e:
            logger.error(f"宽表转长表失败: {e}")
            return pd.DataFrame()

    def _map_english_to_chinese(self, english_name: str) -> str:
        """将英文科目名映射为中文科目名"""
        # 科目名称映射字典
        mapping = {
            'MONETARYFUNDS': '货币资金',
            'ACCOUNTS_RECE': '应收账款',
            'NOTE_RECE': '应收票据',
            'NOTE_ACCOUNTS_RECE': '应收票据及应收账款',
            'ADVANCE_RECEIVABLES': '应收款项融资',
            'PREPAYMENT': '预付款项',
            'OTHER_RECE': '其他应收款',
            'INVENTORY': '存货',
            'CONTRACT_ASSET': '合同资产',
            'HOLDSALE_ASSET': '持有待售资产',
            'NONCURRENT_ASSET_1YEAR': '一年内到期的非流动资产',
            'OTHER_CURRENT_ASSET': '其他流动资产',
            'TOTAL_CURRENT_ASSETS': '流动资产合计',

            'LONG_RECE': '长期应收款',
            'LONG_EQUITY_INVEST': '长期股权投资',
            'OTHER_EQUITY_INVEST': '其他权益工具投资',
            'OTHER_NONCURRENT_FINASSET': '其他非流动金融资产',
            'INVEST_REALESTATE': '投资性房地产',
            'FIXED_ASSET': '固定资产',
            'CIP': '在建工程',
            'USERIGHT_ASSET': '使用权资产',
            'INTANGIBLE_ASSET': '无形资产',
            'DEVELOP_EXPENSE': '开发支出',
            'GOODWILL': '商誉',
            'LONG_PREPAID_EXPENSE': '长期待摊费用',
            'DEFER_TAX_ASSET': '递延所得税资产',
            'OTHER_NONCURRENT_ASSET': '其他非流动资产',
            'TOTAL_NONCURRENT_ASSETS': '非流动资产合计',
            'TOTAL_ASSETS': '资产总计',

            'SHORT_LOAN': '短期借款',
            'BORROW_FUND': '向中央银行借款',
            'LEND_FUND': '拆入资金',
            'TRADE_FINLIAB': '交易性金融负债',
            'DERIVE_FINLIAB': '衍生金融负债',
            'NOTE_PAYABLE': '应付票据',
            'ACCOUNTS_PAYABLE': '应付账款',
            'NOTE_ACCOUNTS_PAYABLE': '应付票据及应付账款',
            'ADVANCE_RECEIVABLES': '预收款项',
            'CONTRACT_LIAB': '合同负债',
            'STAFF_SALARY_PAYABLE': '应付职工薪酬',
            'TAX_PAYABLE': '应交税费',
            'OTHER_PAYABLE': '其他应付款',
            'NONCURRENT_LIAB_1YEAR': '一年内到期的非流动负债',
            'OTHER_CURRENT_LIAB': '其他流动负债',
            'TOTAL_CURRENT_LIAB': '流动负债合计',

            'LONG_LOAN': '长期借款',
            'BOND_PAYABLE': '应付债券',
            'LEASE_LIAB': '租赁负债',
            'LONG_PAYABLE': '长期应付款',
            'LONG_STAFFSALARY_PAYABLE': '长期应付职工薪酬',
            'PREDICT_LIAB': '预计负债',
            'DEFER_INCOME': '递延收益',
            'DEFER_TAX_LIAB': '递延所得税负债',
            'OTHER_NONCURRENT_LIAB': '其他非流动负债',
            'TOTAL_NONCURRENT_LIAB': '非流动负债合计',
            'TOTAL_LIABILITIES': '负债合计',

            'SHARE_CAPITAL': '股本',
            'CAPITAL_RESERVE': '资本公积',
            'TREASURY_SHARES': '减：库存股',
            'OTHER_COMPRE_INCOME': '其他综合收益',
            'SPECIAL_RESERVE': '专项储备',
            'SURPLUS_RESERVE': '盈余公积',
            'UNASSIGN_RPOFIT': '未分配利润',
            'TOTAL_PARENT_EQUITY': '归属于母公司所有者权益合计',
            'MINORITY_EQUITY': '少数股东权益',
            'TOTAL_EQUITY': '所有者权益合计',
            'TOTAL_LIAB_EQUITY': '负债和所有者权益总计',
        }

        return mapping.get(english_name, english_name)

    def _save_to_cache(self, stock_code: str, data: pd.DataFrame):
        """保存数据到缓存"""
        try:
            # 保存股票基本信息
            db_manager.save_stock_info(stock_code)
            
            # 保存资产负债表数据
            db_manager.save_balance_sheet_data(stock_code, data)
            
            logger.info(f"数据已保存到缓存: {stock_code}")
            
        except Exception as e:
            logger.error(f"保存缓存失败: {e}")
    
    def get_stock_list(self) -> pd.DataFrame:
        """获取已缓存的股票列表"""
        try:
            with db_manager.db_path as conn:
                query = """
                    SELECT stock_code, stock_name, market, industry, updated_at
                    FROM stocks
                    ORDER BY updated_at DESC
                """
                return pd.read_sql_query(query, conn)
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return pd.DataFrame()
    
    def clear_cache(self, stock_code: str = None):
        """清理缓存数据"""
        try:
            import sqlite3
            with sqlite3.connect(db_manager.db_path) as conn:
                if stock_code:
                    # 清理特定股票的缓存
                    conn.execute("DELETE FROM balance_sheet WHERE stock_code = ?", (stock_code,))
                    conn.execute("DELETE FROM analysis_cache WHERE stock_code = ?", (stock_code,))
                    logger.info(f"已清理股票 {stock_code} 的缓存")
                else:
                    # 清理所有缓存
                    conn.execute("DELETE FROM balance_sheet")
                    conn.execute("DELETE FROM analysis_cache")
                    logger.info("已清理所有缓存数据")
                
                conn.commit()
                
        except Exception as e:
            logger.error(f"清理缓存失败: {e}")


# 全局数据获取器实例
data_fetcher = DataFetcher()
